{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cd032bcb-fefb-48ec-94da-08d49ac26120",
   "metadata": {},
   "source": [
    "# Query Pipeline with Routing\n",
    "\n",
    "Here we showcase our query pipeline with routing.\n",
    "\n",
    "Routing lets us dynamically choose underlying query pipelines to use given the query and a set of choices.\n",
    "\n",
    "We offer this as an out-of-the-box abstraction in our [Router Query Engine](https://docs.llamaindex.ai/en/stable/examples/query_engine/RouterQueryEngine.html) guide. Here we show you how to compose a similar pipeline using our Query Pipeline syntax - this allows you to not only define query engines but easily stitch it into a chain/DAG with other modules across the compute graph."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3531eedc-4f65-457e-8844-55fcc1773154",
   "metadata": {},
   "source": [
    "## Load Data\n",
    "\n",
    "Load in the Paul Graham essay as an example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6acb8e71",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-llms-openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a441905-9007-44d6-b71a-6fc3e5023e49",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2024-01-10 12:31:00--  https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt\n",
      "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.111.133, 185.199.110.133, 185.199.108.133, ...\n",
      "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.111.133|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 75042 (73K) [text/plain]\n",
      "Saving to: ‘pg_essay.txt’\n",
      "\n",
      "pg_essay.txt        100%[===================>]  73.28K  --.-KB/s    in 0.01s   \n",
      "\n",
      "2024-01-10 12:31:00 (6.32 MB/s) - ‘pg_essay.txt’ saved [75042/75042]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt' -O pg_essay.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3533149c-4312-4444-9b45-52afe21731ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import SimpleDirectoryReader\n",
    "\n",
    "reader = SimpleDirectoryReader(input_files=[\"pg_essay.txt\"])\n",
    "documents = reader.load_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c1d5ff8-ae04-4ea3-bbe0-2c097af71efd",
   "metadata": {},
   "source": [
    "## Setup Query Pipeline with Routing"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63caf998-0a88-4c50-b6a4-2a0c412bde5b",
   "metadata": {},
   "source": [
    "### Define Modules\n",
    "\n",
    "We define llm, vector index, summary index, and prompt templates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01fcbdb2-6747-4e65-b1ce-5d40febccb81",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.query_pipeline import QueryPipeline, InputComponent\n",
    "from typing import Dict, Any, List, Optional\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.core import Document, VectorStoreIndex\n",
    "from llama_index.core import SummaryIndex\n",
    "from llama_index.core.response_synthesizers import TreeSummarize\n",
    "from llama_index.core.schema import NodeWithScore, TextNode\n",
    "from llama_index.core import PromptTemplate\n",
    "from llama_index.core.selectors import LLMSingleSelector\n",
    "\n",
    "# define HyDE template\n",
    "hyde_str = \"\"\"\\\n",
    "Please write a passage to answer the question: {query_str}\n",
    "\n",
    "Try to include as many key details as possible.\n",
    "\n",
    "Passage: \"\"\"\n",
    "hyde_prompt = PromptTemplate(hyde_str)\n",
    "\n",
    "# define llm\n",
    "llm = OpenAI(model=\"gpt-3.5-turbo\")\n",
    "\n",
    "\n",
    "# define synthesizer\n",
    "summarizer = TreeSummarize(llm=llm)\n",
    "\n",
    "# define vector retriever\n",
    "vector_index = VectorStoreIndex.from_documents(documents)\n",
    "vector_query_engine = vector_index.as_query_engine(similarity_top_k=2)\n",
    "\n",
    "# define summary query prompts + retrievers\n",
    "summary_index = SummaryIndex.from_documents(documents)\n",
    "summary_qrewrite_str = \"\"\"\\\n",
    "Here's a question:\n",
    "{query_str}\n",
    "\n",
    "You are responsible for feeding the question to an agent that given context will try to answer the question.\n",
    "The context may or may not be relevant. Rewrite the question to highlight the fact that\n",
    "only some pieces of context (or none) maybe be relevant.\n",
    "\"\"\"\n",
    "summary_qrewrite_prompt = PromptTemplate(summary_qrewrite_str)\n",
    "summary_query_engine = summary_index.as_query_engine()\n",
    "\n",
    "# define selector\n",
    "selector = LLMSingleSelector.from_defaults()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a87a439-88e6-4130-b28f-45268330d3e4",
   "metadata": {},
   "source": [
    "### Construct Query Pipelines\n",
    "\n",
    "Define a query pipeline for vector index, summary index, and join it together with a router."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff95be2e-517f-4632-a7b8-a2e0dec11d73",
   "metadata": {},
   "outputs": [],
   "source": [
    "# define summary query pipeline\n",
    "from llama_index.core.query_pipeline import RouterComponent\n",
    "\n",
    "vector_chain = QueryPipeline(chain=[vector_query_engine])\n",
    "summary_chain = QueryPipeline(\n",
    "    chain=[summary_qrewrite_prompt, llm, summary_query_engine], verbose=True\n",
    ")\n",
    "\n",
    "choices = [\n",
    "    \"This tool answers specific questions about the document (not summary questions across the document)\",\n",
    "    \"This tool answers summary questions about the document (not specific questions)\",\n",
    "]\n",
    "\n",
    "router_c = RouterComponent(\n",
    "    selector=selector,\n",
    "    choices=choices,\n",
    "    components=[vector_chain, summary_chain],\n",
    "    verbose=True,\n",
    ")\n",
    "# top-level pipeline\n",
    "qp = QueryPipeline(chain=[router_c], verbose=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bda05274-09c5-4b56-b2ba-57f445346e73",
   "metadata": {},
   "source": [
    "## Try out Queries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b36ba65-636f-4fe9-8dee-e318cfe9a50d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;38;2;155;135;227m> Running module c0a87442-3165-443d-9709-960e6ddafe7f with input: \n",
      "query: What did the author do during his time in YC?\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;5;200mSelecting component 0: The author used a tool to answer specific questions about the document, which suggests that he was engaged in analyzing and extracting specific information from the document during his time in YC..\n",
      "\u001b[0mDuring his time in YC, the author worked on various tasks related to running Y Combinator. This included selecting and helping founders, dealing with disputes between cofounders, figuring out when people were lying, and fighting with people who maltreated the startups. The author also worked on writing essays and internal software for YC.\n"
     ]
    }
   ],
   "source": [
    "# compare with sync method\n",
    "response = qp.run(\"What did the author do during his time in YC?\")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1cd595f-31d0-4490-8f18-6132ac240c9c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;38;2;155;135;227m> Running module c0a87442-3165-443d-9709-960e6ddafe7f with input: \n",
      "query: What is a summary of this document?\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;5;200mSelecting component 1: The summary questions about the document are answered by this tool..\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module 0e7e9d49-4c92-45a9-b3bf-0e6ab76b51f9 with input: \n",
      "query_str: What is a summary of this document?\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module b0ece4e3-e6cd-4229-8663-b0cd0638683c with input: \n",
      "messages: Here's a question:\n",
      "What is a summary of this document?\n",
      "\n",
      "You are responsible for feeding the question to an agent that given context will try to answer the question.\n",
      "The context may or may not be relev...\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module f247ae78-a71c-4347-ba49-d9357ee93636 with input: \n",
      "input: assistant: What is the summary of the document?\n",
      "\n",
      "\u001b[0mThe document discusses the development and evolution of Lisp as a programming language. It highlights how Lisp was originally created as a formal model of computation and later transformed into a programming language with the assistance of Steve Russell. The document also emphasizes the unique power and elegance of Lisp in comparison to other languages.\n"
     ]
    }
   ],
   "source": [
    "response = qp.run(\"What is a summary of this document?\")\n",
    "print(str(response))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama_index_v2",
   "language": "python",
   "name": "llama_index_v2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
